{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In google colab, \n",
    "# For master version of catalyst, uncomment:\n",
    "# (master version should be fully compatible with this notebook)\n",
    "# ! pip install git+git://github.com/catalyst-team/catalyst.git\n",
    "\n",
    "# For last release version of catalyst, uncomment:\n",
    "# ! pip install catalyst\n",
    "\n",
    "# For specific commit version of catalyst, uncomment:\n",
    "# ! pip install git+http://github.com/catalyst-team/catalyst.git@{commit_hash}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "bs = 32\n",
    "num_workers = 4\n",
    "\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
    "\n",
    "\n",
    "loaders = collections.OrderedDict()\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(\n",
    "    root='./data', train=True,\n",
    "    download=True, transform=data_transform)\n",
    "trainloader = torch.utils.data.DataLoader(\n",
    "    trainset, batch_size=bs,\n",
    "    shuffle=True, num_workers=num_workers)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(\n",
    "    root='./data', train=False,\n",
    "    download=True, transform=data_transform)\n",
    "testloader = torch.utils.data.DataLoader(\n",
    "    testset, batch_size=bs,\n",
    "    shuffle=False, num_workers=num_workers)\n",
    "\n",
    "loaders[\"train\"] = trainloader\n",
    "loaders[\"valid\"] = testloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 5 * 5)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Intro\n",
    "@TODO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "# for graphs use `tensorboard --logdir=./logs`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_EPOCHS = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 1 - typical training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {},
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from catalyst.dl.runner import SupervisedRunner\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_1\"\n",
    "\n",
    "# model, criterion, optimizer\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {},
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# you can use plotly and tensorboard to plot metrics inside jupyter\n",
    "# by default it only plots loss\n",
    "utils.plot_metrics(logdir=logdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 2 - training with scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_2\"\n",
    "\n",
    "# model, criterion, optimizer\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# any Pytorch scheduler supported\n",
    "# scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[3, 8], gamma=0.3)\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=2)\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 3 - training with early stop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "from catalyst.dl.callbacks import EarlyStoppingCallback\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_3\"\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[3, 8], gamma=0.3)\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        EarlyStoppingCallback(patience=2, min_delta=0.01)\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "utils.plot_metrics(logdir=logdir, metrics=[\"loss\", \"_base/lr\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 4 - training with additional metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {},
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "from catalyst.dl.callbacks import EarlyStoppingCallback, AccuracyCallback\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_4\"\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[3, 8], gamma=0.3)\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        AccuracyCallback(accuracy_args=[1, 3, 5]),\n",
    "        EarlyStoppingCallback(patience=2, min_delta=0.01)\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {},
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "utils.plot_metrics(\n",
    "    logdir=logdir, \n",
    "    metrics=[\"loss\", \"accuracy01\", \"accuracy03\", \"_base/lr\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 5 - training with 1cycle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "from catalyst.dl.callbacks import \\\n",
    "    EarlyStoppingCallback, AccuracyCallback\n",
    "from catalyst.contrib.nn import OneCycleLRWithWarmup\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_5\"\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = OneCycleLRWithWarmup(\n",
    "    optimizer,\n",
    "    num_steps=num_epochs, \n",
    "    lr_range=(0.005, 0.00005),\n",
    "    warmup_steps=2,\n",
    "    momentum_range=(0.85, 0.95)\n",
    ")\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        AccuracyCallback(accuracy_args=[1, 3, 5]),\n",
    "        EarlyStoppingCallback(patience=2, min_delta=0.01),\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {},
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "utils.plot_metrics(\n",
    "    logdir=logdir,\n",
    "    step=\"batch\",\n",
    "    metrics=[\"loss\", \"accuracy01\", \"_base/lr\", \"_base/momentum\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup 6 - training without validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "from catalyst.dl.callbacks import EarlyStoppingCallback, AccuracyCallback\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_6\"\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[3, 8], gamma=0.3)\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders={\"train\": loaders[\"train\"]},\n",
    "    valid_loader=\"train\",\n",
    "    callbacks=[\n",
    "        AccuracyCallback(accuracy_args=[1, 3, 5]),\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_metrics(logdir=logdir, step=\"epoch\", metrics=[\"loss\", \"accuracy01\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 7 - pipeline check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_7\"\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[3, 8], gamma=0.3)\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True,\n",
    "    check=True  # here is the trick\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 8 - multi-stage training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "from catalyst.dl.callbacks import EarlyStoppingCallback, AccuracyCallback\n",
    "\n",
    "# experiment setup\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_8\"\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[3, 8], gamma=0.3)\n",
    "\n",
    "# model runner\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# model training - 1\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        AccuracyCallback(accuracy_args=[1, 3, 5]),\n",
    "        EarlyStoppingCallback(patience=2, min_delta=0.01)\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")\n",
    "\n",
    "# model training - 2\n",
    "num_epochs = NUM_EPOCHS\n",
    "logdir = \"./logs/cifar_simple_notebook_8\"\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.05)\n",
    "\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 9 - predict_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "runner_out = runner.predict_loader(\n",
    "    model, loaders[\"valid\"], resume=f\"{logdir}/checkpoints/best.pth\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "runner_out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "# Setup 10 - predict batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "features, targets = next(iter(loaders[\"valid\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "runner_in = {runner.input_key: features}\n",
    "runner_out = runner.predict_batch(runner_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "runner_out[runner.output_key].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}